package com.uraurora.crescent.advanced.tree.impl;

import com.uraurora.crescent.advanced.tree.ITrie;
import com.google.common.collect.Sets;

import java.util.Arrays;
import java.util.Objects;
import java.util.Set;

public class Trie implements ITrie {

    private static final int a_ASCII = 97;

    private final TrieNode root;

    public Trie(){
        this.root = new TrieNode();
    }

    @Override
    public void add(String string) {
        String str = checkString(string);
        TrieNode p = root;
        for (int i = 0; i < str.length(); i++) {
            if(p.next[str.charAt(i) - 'a']==null){
                p.next[str.charAt(i) - 'a'] = new TrieNode();
            }
            p = p.next[str.charAt(i) - 'a'];
            p.num++;
        }
        p.isEnd = true;
    }

    @Override
    public void add(char[] str) {
        add(String.valueOf(str));
    }

    public Trie append(String str){
        this.add(str);
        return this;
    }

    public Trie append(char[] chars){
        this.add(chars);
        return this;
    }

    public Trie append(String[] strs){
        for (String str : strs) {
            add(str);
        }
        return this;
    }

    public Trie append(Iterable<String> strings){
        for (String str:
             strings) {
            this.add(str);
        }
        return this;
    }

    @Override
    public int numByPrefix(String string) {
        TrieNode node = getNodeByPrefix(checkString(string));
        return node == null ? 0 : node.num;
    }

    @Override
    public boolean contains(String str) {
        TrieNode node = getNodeByPrefix(str);
        return node != null && node.isEnd;
    }

    @Override
    public int numByPrefix(char[] str) {
        return numByPrefix(String.valueOf(str));
    }

    /**
     * 获取字典树所有的字符串，使用dfs实现
     * @return 字符串列表
     */
    public Set<String> toSet(){
        Set<String> res = Sets.newHashSet();
        dfsHelper(this.root, new StringBuilder(), res);
        return res;
    }

    /**
     * 根据字符串前缀找到所有以其为前缀的字符串
     * @param str 前缀字符串
     * @return 所包含的字符串集合
     */
    public Set<String> toSetByPrefix(String str){
        Set<String> res = Sets.newHashSet();
        TrieNode root = getNodeByPrefix(str);
        if (root == null){ return res; }
        if (root.isEnd){ res.add(str); }
        dfsHelper(root, new StringBuilder().append(str), res);
        return res;
    }

    private static void dfsHelper(TrieNode root, StringBuilder path, Set<String> res){
        boolean b = Arrays.stream(root.next).allMatch(Objects::isNull);
        if (b){
            res.add(path.toString());
            return;
        }
        for (int i = 0; i < 26; i++) {
            TrieNode node = root.next[i];
            if (node != null){
                path.append((char) (a_ASCII + i));
                if (node.isEnd){ res.add(path.toString()); }
                dfsHelper(node, path, res);
                path.setLength(path.length() - 1);
            }
        }
    }

    static class TrieNode{
        private boolean isEnd;
        private int num;
        private TrieNode[] next = new TrieNode[26];

        TrieNode(){
            for (int i = 0; i < 26; i++) {
                next[i] = null;
            }
            this.num = 0;
            this.isEnd = false;
        }
    }

    private static String checkString(String str){
        if (!Objects.requireNonNull(str).chars().allMatch(Character::isLetter)) {
            throw new IllegalArgumentException("wrong args in trie, u need a full letter str");
        }
        return str.toLowerCase();
    }

    private TrieNode getNodeByPrefix(String str){
        TrieNode p = root;
        for (int i = 0; i < str.length(); i++) {
            if(p.next[str.charAt(i) - 'a']==null){
                return null;
            }
            p = p.next[str.charAt(i) - 'a'];
        }
        return p;
    }


}
